import torch
from torch import nn


class Encoder(nn.Module):
    '''
        编码器：3个全连接
    '''

    def __init__(self):
        super().__init__()
        self.en = nn.Linear(28 * 28 + 10, 400)
        self.mean = nn.Linear(400, 10)
        self.std = nn.Linear(400, 10)
        self.act = nn.ReLU()

    def forward(self, x, l):
        x = torch.cat([x, l], 1)
        x = self.en(x)
        x = self.act(x)
        mean = self.mean(x)
        std = self.std(x)
        return mean, torch.exp(std)


class Decoder(nn.Module):
    '''
        解码器：3个全连接层，未使用转置卷积
    '''

    def __init__(self):
        super().__init__()
        self.de1 = nn.Linear(10 + 10, 400)
        self.de2 = nn.Linear(400, 28 * 28)
        self.act1 = nn.ReLU()
        self.act2 = nn.Sigmoid()

    def forward(self, h, l):
        x = torch.cat([h, l], 1)
        x = self.de1(x)
        x = self.act1(x)
        x = self.de2(x)
        x = self.act2(x)
        return x


class CVAE(nn.Module):
    '''
        CVAE模型构建
    '''

    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def forward(self, x, l):
        '''
            返回encoder和decoder输出
        '''
        mean, std = self.encoder(x, l)
        z = self.SampleMultiGaussian(mean, std)
        x = self.decoder(z, l)
        return x, mean, std

    def SampleMultiGaussian(self, mean, std):
        '''
            利用重参数技巧进行采样
        '''
        device = self.encoder.en._parameters["weight"].device
        e = torch.randn(mean.shape, requires_grad=False).to(device)
        return mean + std * e

    def vae_loss(self, x, x_, mean, std):
        mse = nn.MSELoss(reduction='sum')
        loss1 = mse(x, x_)
        loss2 = -0.5 * (1 + 2 * torch.log(std) - torch.pow(mean, 2) -
                        torch.pow(std, 2)).sum()
        return loss1 + loss2
